import tushare as ts
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.metrics import roc_curve, auc

# 设置 tushare token
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 获取股票的日线行情数据
stock_list = ['000001.SZ', '000002.SZ', '000004.SZ', '000005.SZ', '000006.SZ',
              '000007.SZ', '000008.SZ', '000009.SZ', '000010.SZ', '000011.SZ']
start_date = '20230101'
end_date = '20250101'
data_list = []
for stock in stock_list:
    df = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
    print(f"Stock: {stock}, Columns: {df.columns}")  # 打印列名查看
    data_list.append(df)
all_data = pd.concat(data_list, ignore_index=True)
all_data = all_data.sort_values(by=['ts_code', 'trade_date'])

# 对股票行情数据进行分类（打标签）与技术指标计算
all_data['label'] = all_data.groupby('ts_code')['close'].shift(-1) > all_data['close']
all_data['label'] = all_data['label'].astype(int)

# 计算技术指标
def calculate_technical_indicators(df):
    # 简单移动平均线
    df['SMA_5'] = df['close'].rolling(window=5).mean()
    df['SMA_10'] = df['close'].rolling(window=10).mean()
    df['SMA_20'] = df['close'].rolling(window=20).mean()

    # 指数移动平均线
    df['EMA_5'] = df['close'].ewm(span=5, adjust=False).mean()
    df['EMA_10'] = df['close'].ewm(span=10, adjust=False).mean()
    df['EMA_20'] = df['close'].ewm(span=20, adjust=False).mean()
    df['EMA_12'] = df['close'].ewm(span=12, adjust=False).mean()  # 添加 EMA_12 计算
    df['EMA_26'] = df['close'].ewm(span=26, adjust=False).mean()  # 添加 EMA_26 计算

    # 相对强弱指数
    delta = df['close'].diff()
    up = delta.clip(lower=0)
    down = -delta.clip(upper=0)
    avg_gain = up.rolling(window=14).mean()
    avg_loss = down.rolling(window=14).mean()
    rs = avg_gain / avg_loss
    df['RSI'] = 100 - (100 / (1 + rs))

    # 布林带
    df['STD_20'] = df['close'].rolling(window=20).std()
    df['Upper_Band'] = df['SMA_20'] + (df['STD_20'] * 2)
    df['Lower_Band'] = df['SMA_20'] - (df['STD_20'] * 2)

    # 移动平均收敛背离指标
    df['MACD'] = df['EMA_12'] - df['EMA_26']
    df['Signal_Line'] = df['MACD'].ewm(span=9, adjust=False).mean()

    # 随机指标
    high_14 = df['high'].rolling(window=14).max()
    low_14 = df['low'].rolling(window=14).min()
    df['%K'] = 100 * ((df['close'] - low_14) / (high_14 - low_14))
    df['%D'] = df['%K'].rolling(window=3).mean()

    # 威廉指标
    df['Williams %R'] = -100 * ((high_14 - df['close']) / (high_14 - low_14))

    # 平均趋向指标
    tr = pd.DataFrame()
    tr['H-L'] = df['high'] - df['low']
    tr['H-PC'] = abs(df['high'] - df['close'].shift())
    tr['L-PC'] = abs(df['low'] - df['close'].shift())
    df['TR'] = tr[['H-L', 'H-PC', 'L-PC']].max(axis=1)
    df['ATR'] = df['TR'].rolling(window=14).mean()
    up_move = df['high'] - df['high'].shift()
    down_move = df['low'].shift() - df['low']
    df['+DM'] = np.where((up_move > down_move) & (up_move > 0), up_move, 0)
    df['-DM'] = np.where((down_move > up_move) & (down_move > 0), down_move, 0)
    df['+DI'] = 100 * (df['+DM'].ewm(span=14, adjust=False).mean() / df['ATR'])
    df['-DI'] = 100 * (df['-DM'].ewm(span=14, adjust=False).mean() / df['ATR'])
    df['ADX'] = 100 * (abs(df['+DI'] - df['-DI']) / (df['+DI'] + df['-DI'])).ewm(span=14, adjust=False).mean()

    return df


all_data = all_data.groupby('ts_code').apply(calculate_technical_indicators)

# 对分类计算后的数据进行建模前的处理与分析
# 处理空值
all_data = all_data.dropna()

# 处理异常值
Q1 = all_data.quantile(0.25)
Q3 = all_data.quantile(0.75)
IQR = Q3 - Q1
all_data = all_data[~((all_data < (Q1 - 1.5 * IQR)) | (all_data > (Q3 + 1.5 * IQR))).any(axis=1)]

# 归一化
scaler = MinMaxScaler()
numerical_columns = all_data.select_dtypes(include=[np.number]).columns
all_data[numerical_columns] = scaler.fit_transform(all_data[numerical_columns])

# 主成分分析
pca = PCA(n_components=0.95)
X = all_data.drop(['ts_code', 'trade_date', 'label'], axis=1)
X_pca = pca.fit_transform(X)

# 相关性分析
correlation_matrix = all_data[numerical_columns].corr()
plt.figure(figsize=(12, 10))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm')
plt.title('Correlation Matrix')
plt.show()

# 降维
# 使用 PCA 进行降维
X_reduced = X_pca

# 个股画像
# 用 amount 列替换 volume 列
stock_portraits = all_data.groupby('ts_code')[['close', 'amount', 'SMA_5', 'RSI', 'MACD', 'ADX']].mean()

# 数据均衡（欠采样）
X = all_data.drop(['ts_code', 'trade_date', 'label'], axis=1)
y = all_data['label']
# 统计各类别的数量
class_0_indices = y[y == 0].index
class_1_indices = y[y == 1].index
n_class_0 = len(class_0_indices)
n_class_1 = len(class_1_indices)

# 确定少数类和多数类
if n_class_0 < n_class_1:
    minority_class_indices = class_0_indices
    majority_class_indices = class_1_indices
else:
    minority_class_indices = class_1_indices
    majority_class_indices = class_0_indices

# 从多数类中随机选择与少数类数量相同的样本
random_majority_indices = np.random.choice(majority_class_indices, size=len(minority_class_indices), replace=False)
balanced_indices = np.concatenate([minority_class_indices, random_majority_indices])

X_resampled = X.loc[balanced_indices]
y_resampled = y.loc[balanced_indices]

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_resampled, y_resampled, test_size=0.2, random_state=42)

# 建模
model = LogisticRegression()
model.fit(X_train, y_train)

# 模型评价
y_pred = model.predict(X_test)
print(classification_report(y_test, y_pred))

# 混淆矩阵可视化
cm = confusion_matrix(y_test, y_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.title('Confusion Matrix')
plt.xlabel('Predicted Label')
plt.ylabel('True Label')
plt.show()

# ROC 曲线可视化
y_score = model.predict_proba(X_test)[:, 1]
fpr, tpr, thresholds = roc_curve(y_test, y_score)
roc_auc = auc(fpr, tpr)
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, color='darkorange', lw=2, label='ROC curve (area = %0.2f)' % roc_auc)
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver operating characteristic example')
plt.legend(loc="lower right")
plt.show()